from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import requests
import os
from search import WebSearcher
import embedding_handle


"""
通过 RAG 技术，将外部实时数据（例如网页内容、新闻、文档等）检索并注入到模型的上下文中，使其基于最新信息生成回答

"""


# 初始化 Ollama 模型
llm = Ollama(model="qwen2.5:7b", base_url="http://127.0.0.1:11434")

searcher = WebSearcher()

# 定义联网检索函数
def search_web(query):
    response = searcher.search_web(query=query, num_results=10)
    return  embedding_handle.rank_search_results(query, response.get('results', []))

# 定义提示模板
template = """基于以下最新信息回答问题：
{context}

问题：{question}
"""
prompt = PromptTemplate(template=template, input_variables=["context", "question"])

# 构建问答函数
def ask_with_web(query):
    # 检索实时数据
    context = search_web(query)
    print(f"搜索到的内容摘要：{context}")
    # 将上下文和问题传递给模型
    full_prompt = prompt.format(context=context, question=query)
    response = llm.invoke(full_prompt)
    return response

# 测试
question = "美国最新的关税政策有哪些"
answer = ask_with_web(question)
print(f"----------------------------------------------------------------------------------------------\n\n回答：{answer}")